#include"vld_imu_slam/imu_handle.h"

namespace vld_imu_slam{

void IMUHandle::pushbackIMUData(sensor_msgs::Imu imu_data_converted){
    {
        lock_guard<mutex> lock_integ(imu_integ_mutex_);
        imu_deq_integ_.push_back(imu_data_converted);
    }
    {
        lock_guard<mutex> lock_opt(imu_opt_mutex_);
        imu_deq_opt_.push_back(imu_data_converted);
    }
}

void IMUHandle::setFrameTimestampInteg(const double time){
    unique_lock<mutex> lck(imuInteg_thread_mutex_);
    integ_imu_cur_timestamp_ = time;
    is_set_frametime_integ_.store(true);
    imu_integ_cv_.notify_one();
}

deque<nav_msgs::Odometry> IMUHandle::imuIntegration(){
    deque<nav_msgs::Odometry> deq_odom;
    if (doneFirstOpt_ == false)
        return deq_odom;

    unique_lock<mutex> lck(imuInteg_thread_mutex_);
    imu_integ_cv_.wait(lck,[this] { return is_set_frametime_integ_.load(); });
    is_set_frametime_integ_.store(false);

    sensor_msgs::Imu this_imu;
    {
        lock_guard<mutex> lock_integ(imu_integ_mutex_);
        if(imu_deq_integ_.empty())
            return deq_odom;
        while(imu_deq_integ_.front().header.stamp.toSec()<=integ_imu_cur_timestamp_){
            this_imu = imu_deq_integ_.front();
            imu_deq_integ_.pop_front();

            double dt = (last_imu_timestamp_ < 0) ? (1.0 / 500.0) : (this_imu.header.stamp.toSec() - last_imu_timestamp_);
            last_imu_timestamp_ = this_imu.header.stamp.toSec();

            gtsam::NavState current_state;
            {
                lock_guard<mutex> lck(imuIntegrator_mutex_);
                // imu预积分器添加一帧imu数据，注：这个预积分器的起始时刻是上一帧激光里程计时刻
                imuIntegrator_->integrateMeasurement(gtsam::Vector3(this_imu.linear_acceleration.x, this_imu.linear_acceleration.y, this_imu.linear_acceleration.z),
                                                        gtsam::Vector3(this_imu.angular_velocity.x,    this_imu.angular_velocity.y,    this_imu.angular_velocity.z), dt);

                // 用上一帧激光里程计时刻对应的状态、偏置，施加从该时刻开始到当前时刻的imu预计分量，得到当前时刻的状态
                // current_state = imuIntegrator_->predict(prev_state_opt_, prev_bias_opt_);
                current_state = imuIntegrator_->predict(prev_state_odom_, prev_bias_odom_);
                // current_state = imuIntegrator_->predict(prev_state_odom_, prev_bias_odom_);
            }
        
            // 发布imu里程计（转到lidar系，与激光里程计同一个系）
            nav_msgs::Odometry odometry;
            odometry.header.stamp = this_imu.header.stamp;
            odometry.header.frame_id = "map";
            odometry.child_frame_id = "odom_imu";

            // 变换到lidar系
            gtsam::Pose3 imu_integ_lidar_pose = gtsam::Pose3(current_state.quaternion(), current_state.position());

            odometry.pose.pose.position.x = imu_integ_lidar_pose.translation().x();
            odometry.pose.pose.position.y = imu_integ_lidar_pose.translation().y();
            odometry.pose.pose.position.z = imu_integ_lidar_pose.translation().z();
            odometry.pose.pose.orientation.x = imu_integ_lidar_pose.rotation().toQuaternion().x();
            odometry.pose.pose.orientation.y = imu_integ_lidar_pose.rotation().toQuaternion().y();
            odometry.pose.pose.orientation.z = imu_integ_lidar_pose.rotation().toQuaternion().z();
            odometry.pose.pose.orientation.w = imu_integ_lidar_pose.rotation().toQuaternion().w();
            
            odometry.twist.twist.linear.x = current_state.velocity().x();
            odometry.twist.twist.linear.y = current_state.velocity().y();
            odometry.twist.twist.linear.z = current_state.velocity().z();
            odometry.twist.twist.angular.x = this_imu.angular_velocity.x + prev_bias_odom_.gyroscope().x();
            odometry.twist.twist.angular.y = this_imu.angular_velocity.y + prev_bias_odom_.gyroscope().y();
            odometry.twist.twist.angular.z = this_imu.angular_velocity.z + prev_bias_odom_.gyroscope().z();
            deq_odom.push_back(odometry);
        }
        // cout<<"lidar pose "<<prev_state_odom_.matrix()<<" bias: "<<prev_bias_odom_.vector()<<endl;
    } 
    return deq_odom;
}

void IMUHandle::popUsenessData(deque<sensor_msgs::Imu>& imu_deq,double time){
    double time_temp = time;
    while (!imu_deq.empty())
    {
        if (imu_deq.front().header.stamp.toSec() < time_temp)
        {
            time_temp = imu_deq.front().header.stamp.toSec();
            imu_deq.pop_front();
        }
        else
            break;
    }
}

void IMUHandle::setFrameTimestamp(const double time,const SE3& pose){
    unique_lock<mutex> lck(imu_preinteg_mutex_);
    frame_timestamp_cur_ = time;
    pose_current_ = pose;
    is_get_framestamp_.store(true);
    imu_opt_cv_.notify_one();
}

void IMUHandle::imuPreintegration(){
    // 确保imu优化队列中有imu数据进行预积分
    {
        lock_guard<mutex> lck(imu_opt_mutex_);
        if (imu_deq_opt_.empty())
            return;
    }
    unique_lock<mutex> lck(imu_preinteg_mutex_);
    imu_opt_cv_.wait(lck,[this] { return is_get_framestamp_.load(); });
    is_get_framestamp_.store(false);

    ROS_DEBUG_STREAM("imu handle preintegral keyframe time: "<<frame_timestamp_cur_);

    gtsam::Pose3 lidar_pose_cur_ = SE3togtsamPose(pose_current_);

    // 0. 系统初始化，第一帧
    if (systemInitialized_ == false)
    {
        // 重置ISAM2优化器
        resetOptimization();

        // 从imu优化队列中删除当前帧激光里程计时刻之前的imu数据
        popUsenessData(imu_deq_opt_,frame_timestamp_cur_);
        // 添加里程计位姿先验因子
        prev_pose_opt_ = lidar_pose_cur_;
        gtsam::PriorFactor<gtsam::Pose3> priorPose(X(0), prev_pose_opt_, priorPoseNoise_);
        graphFactors_.add(priorPose);
        // 添加速度先验因子
        prev_vel_opt_ = gtsam::Vector3(0, 0, 0);
        gtsam::PriorFactor<gtsam::Vector3> priorVel(V(0), prev_vel_opt_, priorVelNoise_);
        graphFactors_.add(priorVel);
        // 添加imu偏置先验因子
        prev_bias_opt_ = gtsam::imuBias::ConstantBias();
        gtsam::PriorFactor<gtsam::imuBias::ConstantBias> priorBias(B(0), prev_bias_opt_, priorBiasNoise_);
        graphFactors_.add(priorBias);
        // 变量节点赋初值
        graphValues_.insert(X(0), prev_pose_opt_);
        graphValues_.insert(V(0), prev_vel_opt_);
        graphValues_.insert(B(0), prev_bias_opt_);
        // 优化一次
        isam2_.update(graphFactors_, graphValues_);
        graphFactors_.resize(0);
        graphValues_.clear();
        
        // 重置优化之后的偏置
        {
            lock_guard<mutex> lck(imuIntegrator_mutex_);
            imuIntegrator_->resetIntegrationAndSetBias(prev_bias_opt_);
        }
        imuIntegratorOpt_->resetIntegrationAndSetBias(prev_bias_opt_);
        
        systemInitialized_ = true;
        frame_counter_ = 1;
        return;
    }

    // 每隔100帧激光里程计，重置ISAM2优化器，保证优化效率
    if (frame_counter_ == 100)
    {
        // 前一帧的位姿、速度、偏置噪声模型
        gtsam::noiseModel::Gaussian::shared_ptr updatedPoseNoise = gtsam::noiseModel::Gaussian::Covariance(isam2_.marginalCovariance(X(frame_counter_-1)));
        gtsam::noiseModel::Gaussian::shared_ptr updatedVelNoise  = gtsam::noiseModel::Gaussian::Covariance(isam2_.marginalCovariance(V(frame_counter_-1)));
        gtsam::noiseModel::Gaussian::shared_ptr updatedBiasNoise = gtsam::noiseModel::Gaussian::Covariance(isam2_.marginalCovariance(B(frame_counter_-1)));
        // 重置ISAM2优化器
        resetOptimization();
        // 添加位姿先验因子，用前一帧的值初始化
        gtsam::PriorFactor<gtsam::Pose3> priorPose(X(0), prev_pose_opt_, updatedPoseNoise);
        graphFactors_.add(priorPose);
        // 添加速度先验因子，用前一帧的值初始化
        gtsam::PriorFactor<gtsam::Vector3> priorVel(V(0), prev_vel_opt_, updatedVelNoise);
        graphFactors_.add(priorVel);
        // 添加偏置先验因子，用前一帧的值初始化
        gtsam::PriorFactor<gtsam::imuBias::ConstantBias> priorBias(B(0), prev_bias_opt_, updatedBiasNoise);
        graphFactors_.add(priorBias);
        // 变量节点赋初值，用前一帧的值初始化
        graphValues_.insert(X(0), prev_pose_opt_);
        graphValues_.insert(V(0), prev_vel_opt_);
        graphValues_.insert(B(0), prev_bias_opt_);
        // 优化一次
        isam2_.update(graphFactors_, graphValues_);
        graphFactors_.resize(0);
        graphValues_.clear();

        frame_counter_ = 1;
    }


    // 1. 计算前一帧与当前帧之间的imu预积分量，用前一帧状态施加预积分量得到当前帧初始状态估计，添加来自mapOptimization的当前帧位姿，进行因子图优化，更新当前帧状态
    {
        lock_guard<mutex> lck(imu_opt_mutex_);
        while (!imu_deq_opt_.empty())
        {
            // 提取前一帧与当前帧之间的imu数据，计算预积分
            sensor_msgs::Imu *thisImu = &imu_deq_opt_.front();
            if (thisImu->header.stamp.toSec() < frame_timestamp_cur_)
            {
                // 两帧imu数据时间间隔
                double dt = (last_imu_timestamp_opt_ < 0) ? (1.0 / 500.0) : (thisImu->header.stamp.toSec() - last_imu_timestamp_opt_);
                // imu预积分数据输入：加速度、角速度、dt
                imuIntegratorOpt_->integrateMeasurement(
                        gtsam::Vector3(thisImu->linear_acceleration.x, thisImu->linear_acceleration.y, thisImu->linear_acceleration.z),
                        gtsam::Vector3(thisImu->angular_velocity.x,    thisImu->angular_velocity.y,    thisImu->angular_velocity.z), dt);
                
                last_imu_timestamp_opt_ = thisImu->header.stamp.toSec();
                // 从队列中删除已经处理的imu数据
                imu_deq_opt_.pop_front();
            }
            else
                break;
        }
    }

    // 添加imu预积分因子
    const gtsam::PreintegratedImuMeasurements& preint_imu = dynamic_cast<const gtsam::PreintegratedImuMeasurements&>(*imuIntegratorOpt_);
    // 参数：前一帧位姿，前一帧速度，当前帧位姿，当前帧速度，前一帧偏置，预计分量
    gtsam::ImuFactor imu_factor(X(frame_counter_ - 1), V(frame_counter_ - 1), X(frame_counter_), V(frame_counter_), B(frame_counter_ - 1), preint_imu);
    graphFactors_.add(imu_factor);
    // 添加imu偏置因子，前一帧偏置，当前帧偏置，观测值，噪声协方差；deltaTij()是积分段的时间
    graphFactors_.add(gtsam::BetweenFactor<gtsam::imuBias::ConstantBias>(B(frame_counter_ - 1), B(frame_counter_), gtsam::imuBias::ConstantBias(),
                        gtsam::noiseModel::Diagonal::Sigmas(sqrt(imuIntegratorOpt_->deltaTij()) * noiseModelBetweenBias_)));
    // 添加位姿因子
    gtsam::PriorFactor<gtsam::Pose3> pose_factor(X(frame_counter_),lidar_pose_cur_ ,  correctionNoise_);
    graphFactors_.add(pose_factor);
    // 用前一帧的状态、偏置，施加imu预计分量，得到当前帧的状态
    gtsam::NavState propState_ = imuIntegratorOpt_->predict(prev_state_opt_, prev_bias_opt_);
    // 变量节点赋初值
    graphValues_.insert(X(frame_counter_), propState_.pose());
    graphValues_.insert(V(frame_counter_), propState_.v());
    graphValues_.insert(B(frame_counter_), prev_bias_opt_);
    // 优化
    isam2_.update(graphFactors_, graphValues_);
    isam2_.update();
    graphFactors_.resize(0);
    graphValues_.clear();
    // 优化结果
    gtsam::Values result = isam2_.calculateEstimate();
    // 更新当前帧位姿、速度
    prev_pose_opt_  = result.at<gtsam::Pose3>(X(frame_counter_));
    prev_vel_opt_   = result.at<gtsam::Vector3>(V(frame_counter_));
    // 更新当前帧状态
    prev_state_opt_ = gtsam::NavState(prev_pose_opt_, prev_vel_opt_);
    // 更新当前帧imu偏置
    prev_bias_opt_  = result.at<gtsam::imuBias::ConstantBias>(B(frame_counter_));
    // 重置预积分器，设置新的偏置，这样下一帧激光里程计进来的时候，预积分量就是两帧之间的增量
    imuIntegratorOpt_->resetIntegrationAndSetBias(prev_bias_opt_);

    ROS_DEBUG_STREAM("imu hadle imu bias: "<<prev_bias_opt_.vector()<<" velocity: "<<prev_vel_opt_.matrix()<<" pose: "<<prev_pose_opt_.matrix());
    // cout<<"lidar pose "<<lidar_pose_cur_.matrix()<<endl;
    // imu因子图优化结果，速度或者偏置过大，认为失败
    if (failureDetection(prev_vel_opt_, prev_bias_opt_))
    {
        // 重置参数
        resetParams();
        return;
    }


    // 2. 优化之后，执行重传播；优化更新了imu的偏置，用最新的偏置重新计算当前激光里程计时刻之后的imu预积分，这个预积分用于计算每时刻位姿
    {
        lock_guard<mutex> lck(imuIntegrator_mutex_);
        prev_state_odom_ = prev_state_opt_;
        prev_bias_odom_  = prev_bias_opt_;
        imuIntegrator_->resetIntegrationAndSetBias(prev_bias_odom_);
    }   
    // 从imu队列中删除当前激光里程计时刻之前的imu数据
    {
        lock_guard<mutex> lck(imu_integ_mutex_);
        popUsenessData(imu_deq_integ_,frame_timestamp_cur_);
    }

    ++frame_counter_;
    doneFirstOpt_ = true;
}

void IMUHandle::resetOptimization()
{
    gtsam::ISAM2Params optParameters;
    optParameters.relinearizeThreshold = 0.1;
    optParameters.relinearizeSkip = 1;
    isam2_ = gtsam::ISAM2(optParameters);

    gtsam::NonlinearFactorGraph newGraphFactors;
    graphFactors_ = newGraphFactors;

    gtsam::Values NewGraphValues;
    graphValues_ = NewGraphValues;
}

bool IMUHandle::failureDetection(const gtsam::Vector3& velCur, const gtsam::imuBias::ConstantBias& biasCur)
{
    Eigen::Vector3f vel(velCur.x(), velCur.y(), velCur.z());
    if (vel.norm() > 30)
    {
        ROS_WARN("Large velocity, reset IMU-preintegration!");
        return true;
    }

    Eigen::Vector3f ba(biasCur.accelerometer().x(), biasCur.accelerometer().y(), biasCur.accelerometer().z());
    Eigen::Vector3f bg(biasCur.gyroscope().x(), biasCur.gyroscope().y(), biasCur.gyroscope().z());
    if (ba.norm() > 1.0 || bg.norm() > 1.0)
    {
        ROS_WARN("Large bias, reset IMU-preintegration!");
        return true;
    }

    return false;
}

void IMUHandle::resetParams()
{
    last_imu_timestamp_opt_ = -1;
    doneFirstOpt_ = false;
    systemInitialized_ = false;
}


}
